{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6d4ef9c5",
   "metadata": {},
   "source": [
    "$$ y = \\frac{x-E[x]}{\\sqrt{Var[x]+\\epsilon}}*\\gamma + \\beta $$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e796f1c",
   "metadata": {},
   "source": [
    "$$ x = S*q $$\n",
    "$$ \\gamma = S*q_{\\gamma} $$\n",
    "$$ \\beta = S*q_{\\beta}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0440b1cb",
   "metadata": {},
   "source": [
    "$$ y = \\frac{S*q-E[S*q]}{\\sqrt{Var[S*q]+\\epsilon}}*S*q_{\\gamma} + S*q_{\\beta} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4349f425",
   "metadata": {},
   "source": [
    "$$ y = \\frac{S(q-E[q])}{\\sqrt{S^{2}(Var[q]+\\frac{\\epsilon}{S^2})}}*S*q_{\\gamma} + S*q_{\\beta} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c1141f8",
   "metadata": {},
   "source": [
    "$$ y = \\frac{q-E[q]}{\\sqrt{(Var[q]+\\frac{\\epsilon}{S^2})}}*S*q_{\\gamma} + S*q_{\\beta} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2a3b084",
   "metadata": {},
   "source": [
    "$$ y = S(\\frac{q-E[q]}{\\sqrt{(Var[q]+\\frac{\\epsilon}{S^2})}}*q_{\\gamma} + q_{\\beta}) $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "f07b9ed9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at textattack/roberta-base-MRPC were not used when initializing RobertaForSequenceClassification: ['roberta.pooler.dense.bias', 'roberta.pooler.dense.weight']\n",
      "- This IS expected if you are initializing RobertaForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing RobertaForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import numpy as np\n",
    "import torch\n",
    "from src.quant_roberta import QuantRoberta\n",
    "from src.quant_ops import tensor_quant_layernorm\n",
    "model = QuantRoberta()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "1bce3134",
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = torch.from_numpy(np.random.uniform(size=(32,512,768))).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "fd082ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "layernorm = model.model.roberta.encoder.layer[0].output.LayerNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "e8a5d691",
   "metadata": {},
   "outputs": [],
   "source": [
    "gt = layernorm(inp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "b200e172",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[9357.],\n",
      "         [9578.],\n",
      "         [9175.],\n",
      "         ...,\n",
      "         [9407.],\n",
      "         [9752.],\n",
      "         [9539.]],\n",
      "\n",
      "        [[9233.],\n",
      "         [9199.],\n",
      "         [9785.],\n",
      "         ...,\n",
      "         [9501.],\n",
      "         [9463.],\n",
      "         [9384.]],\n",
      "\n",
      "        [[9508.],\n",
      "         [9480.],\n",
      "         [9362.],\n",
      "         ...,\n",
      "         [9416.],\n",
      "         [9362.],\n",
      "         [9432.]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[9349.],\n",
      "         [9568.],\n",
      "         [9011.],\n",
      "         ...,\n",
      "         [9397.],\n",
      "         [9531.],\n",
      "         [9426.]],\n",
      "\n",
      "        [[9487.],\n",
      "         [9406.],\n",
      "         [9305.],\n",
      "         ...,\n",
      "         [9420.],\n",
      "         [9640.],\n",
      "         [9718.]],\n",
      "\n",
      "        [[9368.],\n",
      "         [9670.],\n",
      "         [9330.],\n",
      "         ...,\n",
      "         [9224.],\n",
      "         [9558.],\n",
      "         [9534.]]])\n"
     ]
    }
   ],
   "source": [
    "test = tensor_quant_layernorm(layernorm, inp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "aa2b5711",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.2719e-09, grad_fn=<MeanBackward0>)"
      ]
     },
     "execution_count": 138,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "((gt - test)**2).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "id": "dfeb93fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[-0.0132, -0.2497,  0.0147,  ..., -0.3016, -0.2678,  0.3225],\n",
      "         [-0.5154, -0.6211,  0.6647,  ..., -0.1092, -0.8612, -1.1108],\n",
      "         [ 0.2995,  0.9043, -0.6428,  ..., -0.4918,  0.6762,  0.0958],\n",
      "         ...,\n",
      "         [-0.2391, -0.7679, -0.1575,  ...,  0.1055, -0.5153, -1.3152],\n",
      "         [-0.6596, -0.1150,  0.2063,  ..., -0.2385, -0.3908, -0.8548],\n",
      "         [-0.7113,  0.7206,  0.2078,  ..., -0.7478,  0.0204, -0.6075]],\n",
      "\n",
      "        [[-0.2663, -0.8407, -0.5942,  ..., -0.1754,  0.8408, -0.1712],\n",
      "         [-0.2235, -0.1463,  0.2708,  ...,  1.0346,  0.1868, -0.5495],\n",
      "         [ 0.0155,  0.3317, -0.2293,  ...,  0.4074, -0.4448, -1.3831],\n",
      "         ...,\n",
      "         [-0.4816,  0.8432,  0.1118,  ...,  0.0288, -0.8741,  0.8347],\n",
      "         [ 0.0215, -0.5429, -0.1779,  ..., -0.9427, -0.7520,  0.7450],\n",
      "         [-0.0605,  1.1092,  0.7791,  ...,  0.2658,  0.6505, -0.6987]],\n",
      "\n",
      "        [[-0.7302,  0.4232,  0.8217,  ...,  0.4624,  0.4853, -1.3755],\n",
      "         [ 0.5133,  0.1992, -0.7373,  ..., -0.5069,  0.4360, -1.2885],\n",
      "         [-0.6096, -0.7512,  0.6586,  ...,  0.6961,  0.6221,  0.4755],\n",
      "         ...,\n",
      "         [-0.4939,  0.1931, -0.8660,  ..., -0.9267,  0.2032, -1.2572],\n",
      "         [ 0.3378,  0.0846,  0.1474,  ...,  0.9274,  0.3271, -0.2950],\n",
      "         [-0.5313, -0.2358, -0.5384,  ..., -0.8432,  0.7538,  0.7794]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[-0.5542, -0.3790,  0.4228,  ..., -0.1574, -0.3145, -0.6201],\n",
      "         [-0.0428,  0.9110,  0.7180,  ..., -0.4408, -0.3769, -0.8046],\n",
      "         [ 0.4530, -0.2857,  0.8317,  ...,  0.8155,  0.6260, -1.1260],\n",
      "         ...,\n",
      "         [ 0.3181,  0.4374, -0.9110,  ...,  0.5903,  0.3087, -0.6159],\n",
      "         [-0.6290,  0.5230, -0.0151,  ..., -0.5641, -1.0135, -0.0538],\n",
      "         [-0.1883,  0.3868,  0.2601,  ...,  0.2802,  0.8630, -0.6236]],\n",
      "\n",
      "        [[ 0.3716, -0.4762, -0.1652,  ...,  0.9353, -0.4015,  0.5028],\n",
      "         [-0.5688,  0.2317, -0.5609,  ...,  1.0511,  0.1934,  0.0639],\n",
      "         [-0.6766, -0.7215, -0.2024,  ..., -0.4398, -0.8316,  0.5158],\n",
      "         ...,\n",
      "         [ 0.4996,  1.0216,  0.2838,  ..., -0.4396,  0.1761, -1.1686],\n",
      "         [ 0.1533,  0.2560, -0.2675,  ...,  1.0912,  0.7688, -0.7959],\n",
      "         [-0.2354, -0.4244,  0.0913,  ...,  0.0050,  0.1298,  0.8810]],\n",
      "\n",
      "        [[ 0.0942,  0.5048,  0.1182,  ...,  0.2176, -0.1427, -0.7736],\n",
      "         [ 0.3200,  1.0602,  0.2598,  ..., -0.5710,  0.7747, -1.1266],\n",
      "         [-0.5528, -0.3697, -0.6424,  ...,  0.6186, -0.2355, -0.4868],\n",
      "         ...,\n",
      "         [-0.3420, -0.8489, -0.1368,  ..., -0.8569, -0.8403, -0.0746],\n",
      "         [-0.7064,  0.9793,  0.2324,  ..., -0.5797, -0.3607,  0.4541],\n",
      "         [ 0.5581, -0.9011, -0.5006,  ...,  0.4867, -0.2755,  0.2571]]])\n"
     ]
    }
   ],
   "source": [
    "print(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "52ead9d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.0132, -0.2496,  0.0147,  ..., -0.3016, -0.2677,  0.3225],\n",
       "         [-0.5154, -0.6211,  0.6647,  ..., -0.1092, -0.8611, -1.1108],\n",
       "         [ 0.2995,  0.9042, -0.6428,  ..., -0.4918,  0.6761,  0.0958],\n",
       "         ...,\n",
       "         [-0.2392, -0.7678, -0.1575,  ...,  0.1055, -0.5153, -1.3152],\n",
       "         [-0.6596, -0.1149,  0.2063,  ..., -0.2385, -0.3907, -0.8547],\n",
       "         [-0.7113,  0.7206,  0.2078,  ..., -0.7478,  0.0205, -0.6075]],\n",
       "\n",
       "        [[-0.2662, -0.8406, -0.5942,  ..., -0.1754,  0.8408, -0.1712],\n",
       "         [-0.2234, -0.1462,  0.2708,  ...,  1.0345,  0.1868, -0.5495],\n",
       "         [ 0.0155,  0.3317, -0.2293,  ...,  0.4073, -0.4448, -1.3831],\n",
       "         ...,\n",
       "         [-0.4815,  0.8431,  0.1118,  ...,  0.0288, -0.8740,  0.8346],\n",
       "         [ 0.0216, -0.5428, -0.1779,  ..., -0.9427, -0.7520,  0.7451],\n",
       "         [-0.0605,  1.1091,  0.7791,  ...,  0.2658,  0.6504, -0.6987]],\n",
       "\n",
       "        [[-0.7301,  0.4232,  0.8216,  ...,  0.4624,  0.4852, -1.3755],\n",
       "         [ 0.5133,  0.1992, -0.7373,  ..., -0.5069,  0.4360, -1.2885],\n",
       "         [-0.6096, -0.7511,  0.6585,  ...,  0.6961,  0.6221,  0.4754],\n",
       "         ...,\n",
       "         [-0.4939,  0.1930, -0.8660,  ..., -0.9267,  0.2031, -1.2572],\n",
       "         [ 0.3378,  0.0846,  0.1475,  ...,  0.9274,  0.3271, -0.2950],\n",
       "         [-0.5313, -0.2357, -0.5384,  ..., -0.8431,  0.7538,  0.7794]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[-0.5541, -0.3789,  0.4228,  ..., -0.1574, -0.3145, -0.6200],\n",
       "         [-0.0428,  0.9109,  0.7180,  ..., -0.4408, -0.3768, -0.8045],\n",
       "         [ 0.4530, -0.2857,  0.8317,  ...,  0.8154,  0.6259, -1.1259],\n",
       "         ...,\n",
       "         [ 0.3181,  0.4374, -0.9109,  ...,  0.5903,  0.3087, -0.6159],\n",
       "         [-0.6289,  0.5230, -0.0151,  ..., -0.5641, -1.0135, -0.0538],\n",
       "         [-0.1883,  0.3868,  0.2601,  ...,  0.2802,  0.8630, -0.6235]],\n",
       "\n",
       "        [[ 0.3716, -0.4761, -0.1652,  ...,  0.9352, -0.4014,  0.5027],\n",
       "         [-0.5687,  0.2317, -0.5609,  ...,  1.0511,  0.1934,  0.0639],\n",
       "         [-0.6765, -0.7214, -0.2024,  ..., -0.4399, -0.8315,  0.5157],\n",
       "         ...,\n",
       "         [ 0.4996,  1.0215,  0.2837,  ..., -0.4396,  0.1760, -1.1686],\n",
       "         [ 0.1533,  0.2560, -0.2674,  ...,  1.0912,  0.7688, -0.7958],\n",
       "         [-0.2354, -0.4244,  0.0913,  ...,  0.0050,  0.1298,  0.8810]],\n",
       "\n",
       "        [[ 0.0942,  0.5047,  0.1182,  ...,  0.2176, -0.1426, -0.7736],\n",
       "         [ 0.3200,  1.0601,  0.2598,  ..., -0.5710,  0.7746, -1.1266],\n",
       "         [-0.5527, -0.3697, -0.6424,  ...,  0.6186, -0.2354, -0.4868],\n",
       "         ...,\n",
       "         [-0.3420, -0.8488, -0.1368,  ..., -0.8569, -0.8402, -0.0746],\n",
       "         [-0.7064,  0.9793,  0.2324,  ..., -0.5797, -0.3607,  0.4541],\n",
       "         [ 0.5580, -0.9010, -0.5005,  ...,  0.4867, -0.2755,  0.2570]]],\n",
       "       grad_fn=<NativeLayerNormBackward0>)"
      ]
     },
     "execution_count": 140,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1db9f9a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "trans-fat",
   "language": "python",
   "name": "trans-fat"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
